import torch
from diffusers import StableDiffusionPipeline
from PIL import Image
from io import BytesIO
from algorithms.generate.inferencer_base_generate import InferencerBaseGenerate


class StableDiffusionInferencer(InferencerBaseGenerate):
    def _load_model(self, model_name):
        """加载 Stable Diffusion 模型"""
        print(f"Loading Stable Diffusion model: {model_name}")
        pipeline = StableDiffusionPipeline.from_pretrained(model_name)
        pipeline.to(self.device)
        pipeline.enable_attention_slicing()  # 减少显存占用
        return pipeline

    def preprocess(self, input_data):
        """无需预处理，因为 Stable Diffusion 直接接受文本输入"""
        return input_data

    def postprocess(self, outputs):
        """将生成的图像转换为 PIL.Image 对象"""
        images = []
        for output in outputs.images:
            image = Image.open(BytesIO(output))
            images.append(image)
        return images

    def inference_txt2img(self, input_txt):
        """根据文本生成图像"""
        with torch.no_grad():
            images = self.model(prompt=input_txt)
        return images

    def save_image(self, image, output_path="output.png"):
        """保存生成的图像"""
        image.save(output_path)
        print(f"Image saved to {output_path}")

    def show_image(self, image):
        """显示生成的图像"""
        image.show()



